import torch
from PIL import Image
from torchvision import transforms
from module import *

cifar10 = ['飞机','小车','鸟','猫','鹿','狗','青蛙','马','船','大车']

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

testmodel = torch.load(r'./weight/0923157.pth')

img = Image.open(r'./test_image/小车.JPG').convert('RGB')
trans =  transforms.Compose([ transforms.Resize(size=(32,32)), transforms.ToTensor()])
img = trans(img)
img = torch.reshape(img,(1,3,32,32))

testmodel.to(device)
img = img.to(device)

testmodel.eval()
with torch.no_grad():
    output = testmodel(img).argmax(1) 
    print(cifar10[output.item()])